from torch.nn import functional as F
import matplotlib.pyplot as plt
from PIL import Image
import random
import torch
from cnn import Net
from train_dataset import transform, mean, std
from fg import dataset_details

# 检测是否有可用的GPU，如果没有则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 实例化模型并加载预训练权重
model = Net().to(device)
# 加载模型时指定weights_only=True以增加安全性
model.load_state_dict(torch.load('Net.pth', map_location=device, weights_only=True))
model.eval()

# 提供了两种测试方法，第一种是从测试集随机取10张进行预测，可视化输出
def test_1():
    with open('test.txt', mode='r', encoding='gbk') as f1:
        image_list = random.sample(f1.readlines(), 10)
    f1.close()

    r = 0
    plt.figure(figsize=(20, 20))
    transform_ = transform(mean, std)
    for _, i in enumerate(image_list):
        img_path, target_ = i.strip().split('\t')[0], i.strip().split('\t')[1]
        img = Image.open(img_path).convert('RGB')
        img_ = transform_(img).unsqueeze(0).to(device)
        out = model(img_)
        pred = torch.argmax(F.softmax(out, dim=1)).to('cpu')
        if pred == int(target_):
            r += 1
        plt.subplot(2, 5, _ + 1)
        plt.title(f'pred:{dataset_details[int(pred)]}, acc:{dataset_details[int(target_)]}', fontsize=15)
        plt.imshow(img)
        plt.xticks([])
        plt.yticks([])
    print(f'抽测样本的准确率为：{(r / 10) * 100}%')
    plt.tight_layout()
    plt.show()

def test_2():
    with open('test.txt', mode='r', encoding='gbk') as f1:
        image_list = f1.readlines()
    f1.close()

    r = 0
    transform_ = transform(mean, std)
    for _, i in enumerate(image_list):
        img_path, target_ = i.strip().split('\t')[0], i.strip().split('\t')[1]
        # 检查路径是否包含 '__MACOSX'，如果是，则跳过
        if '__MACOSX' in img_path:
            continue
        try:
            img = Image.open(img_path).convert('RGB')
            img_ = transform_(img).unsqueeze(0).to(device)
            out = model(img_)
            pred = torch.argmax(F.softmax(out, dim=1)).to('cpu')
            if pred == int(target_):
                r += 1
        except PermissionError:
            print(f"Permission denied: {img_path}")
            continue
        except Exception as e:
            print(f"Error opening {img_path}: {e}")
            continue
    print(f'测试集准确率为：{(r / len(image_list)) * 100}%')

test_1()
#test_2()
